package net.cyue.web.easyquery.provider.db.jdbc;

import net.cyue.web.easyquery.provider.db.DBConfiguration;

import javax.sql.DataSource;
import java.io.PrintWriter;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.sql.Statement;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;

public class SimpleDataSource implements DataSource {
    private final String driverClass;
    private final String url;
    private final String username;
    private final String password;
    private final int initialSize;
    private final int maxSize;
    private final long maxWaitMillis;
    private final long connectionTimeoutMillis;
    private final String validationQuery;

    // 使用线程安全的阻塞队列管理连接池
    private final BlockingQueue<PooledConnection> connectionPool;
    private int activeConnections;

    // 包装Connection，添加创建时间和状态跟踪
    private static class PooledConnection {
        final Connection connection;
        final long createTime;
        boolean isValid;

        PooledConnection(Connection connection) {
            this.connection = connection;
            this.createTime = System.currentTimeMillis();
            this.isValid = true;
        }
    }

    public SimpleDataSource(DBConfiguration configuration) {
        this(
            builder()
            .driverClass(configuration.getDriverName())
            .url(configuration.getUrl())
            .username(configuration.getUsername())
            .password(configuration.getPassword())
        );
    }

    // 构造函数 - 使用建造者模式更灵活地设置参数
    private SimpleDataSource(Builder builder) {
        this.driverClass = builder.driverClass;
        this.url = builder.url;
        this.username = builder.username;
        this.password = builder.password;
        this.initialSize = builder.initialSize;
        this.maxSize = builder.maxSize;
        this.maxWaitMillis = builder.maxWaitMillis;
        this.connectionTimeoutMillis = builder.connectionTimeoutMillis;
        this.validationQuery = builder.validationQuery;

        // 初始化有界阻塞队列，保证线程安全
        this.connectionPool = new ArrayBlockingQueue<>(maxSize);
        this.activeConnections = 0;

        initializeConnectionPool();
    }

    // 初始化连接池
    private void initializeConnectionPool() {
        try {
            // 加载数据库驱动
            Class.forName(driverClass);

            // 预创建初始连接
            for (int i = 0; i < initialSize; i++) {
                PooledConnection pooledConn = createPooledConnection();
                connectionPool.offer(pooledConn);
            }

            // 启动连接清理线程，定期移除过期连接
            startConnectionCleaner();

        } catch (ClassNotFoundException | SQLException e) {
            throw new RuntimeException("Failed to initialize connection pool", e);
        }
    }

    // 创建新的数据库连接并包装
    private PooledConnection createPooledConnection() throws SQLException {
        Connection connection = DriverManager.getConnection(url, username, password);
        // 设置连接不自动提交，由用户控制事务
        connection.setAutoCommit(false);
        return new PooledConnection(connection);
    }

    // 启动连接清理线程
    private void startConnectionCleaner() {
        Thread cleanerThread = new Thread(() -> {
            while (!Thread.currentThread().isInterrupted()) {
                try {
                    // 每30秒检查一次连接
                    TimeUnit.SECONDS.sleep(30);
                    cleanExpiredConnections();
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    break;
                }
            }
        }, "ConnectionCleaner");
        cleanerThread.setDaemon(true); // 守护线程，随主线程退出
        cleanerThread.start();
    }

    // 清理过期连接
    private void cleanExpiredConnections() {
        if (connectionTimeoutMillis <= 0) return;

        long now = System.currentTimeMillis();
        PooledConnection[] connections = connectionPool.toArray(new PooledConnection[0]);

        for (PooledConnection conn : connections) {
            if (now - conn.createTime > connectionTimeoutMillis) {
                // 移除并关闭过期连接
                if (connectionPool.remove(conn)) {
                    closeConnection(conn.connection);
                }
            } else if (!isConnectionValid(conn.connection)) {
                // 移除并关闭无效连接
                if (connectionPool.remove(conn)) {
                    closeConnection(conn.connection);
                }
            }
        }

        // 补充连接到初始数量
        try {
            while (connectionPool.size() < initialSize && activeConnections + connectionPool.size() < maxSize) {
                connectionPool.offer(createPooledConnection());
            }
        } catch (SQLException e) {
            Logger.getLogger(SimpleDataSource.class.getName()).warning("Failed to replenish connections: " + e.getMessage());
        }
    }

    // 验证连接是否有效
    private boolean isConnectionValid(Connection connection) {
        if (validationQuery == null || validationQuery.isEmpty()) {
            try {
                return !connection.isClosed();
            } catch (SQLException e) {
                return false;
            }
        }

        try (Statement stmt = connection.createStatement()) {
            stmt.execute(validationQuery);
            return true;
        } catch (SQLException e) {
            return false;
        }
    }

    // 从连接池获取连接
    @Override
    public Connection getConnection() throws SQLException {
        PooledConnection pooledConn;

        // 尝试从连接池获取连接
        try {
            // 如果队列空了，等待指定时间
            pooledConn = connectionPool.poll(maxWaitMillis, TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new SQLException("Interrupted while waiting for connection", e);
        }

        // 如果没有获取到连接，且未达到最大连接数，则创建新连接
        if (pooledConn == null) {
            synchronized (this) {
                if (activeConnections < maxSize) {
                    pooledConn = createPooledConnection();
                } else {
                    throw new SQLException("Maximum connection limit reached (" + maxSize + ")");
                }
            }
        } else {
            // 验证连接有效性
            if (!isConnectionValid(pooledConn.connection)) {
                closeConnection(pooledConn.connection);
                pooledConn = createPooledConnection();
            }
        }

        // 增加活跃连接计数
        synchronized (this) {
            activeConnections++;
        }

        // 返回代理连接，实现自动归还
        return createProxyConnection(pooledConn);
    }

    // 创建代理连接，当连接关闭时自动归还到连接池
    private Connection createProxyConnection(PooledConnection pooledConn) {
        return new ConnectionProxy(pooledConn.connection) {
            @Override
            public void close() throws SQLException {
                // 当用户调用close()时，将连接归还到池而不是真正关闭
                if (!isClosed()) {
                    super.close(); // 调用代理的close方法
                    returnConnection(pooledConn);
                }
            }
        };
    }

    // 将连接返回连接池
    private void returnConnection(PooledConnection pooledConn) {
        // 减少活跃连接计数
        synchronized (this) {
            activeConnections--;
        }

        // 检查连接是否仍然有效
        if (isConnectionValid(pooledConn.connection)) {
            try {
                // 重置连接状态，准备供下次使用
                if (!pooledConn.connection.getAutoCommit()) {
                    pooledConn.connection.rollback(); // 回滚未提交的事务
                }
                pooledConn.connection.setAutoCommit(false); // 重置自动提交状态
                connectionPool.offer(pooledConn);
            } catch (SQLException e) {
                closeConnection(pooledConn.connection);
                try {
                    connectionPool.offer(createPooledConnection());
                } catch (SQLException ex) {
                    Logger.getLogger(SimpleDataSource.class.getName()).warning("Failed to replace bad connection: " + ex.getMessage());
                }
            }
        } else {
            closeConnection(pooledConn.connection);
            try {
                connectionPool.offer(createPooledConnection());
            } catch (SQLException e) {
                Logger.getLogger(SimpleDataSource.class.getName()).warning("Failed to replace bad connection: " + e.getMessage());
            }
        }
    }

    // 带用户名和密码的getConnection方法
    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        // 这里可以实现基于不同用户的连接管理
        // 简单实现：创建一个新的数据源实例
        return new SimpleDataSource.Builder()
                .driverClass(driverClass)
                .url(url)
                .username(username)
                .password(password)
                .initialSize(initialSize)
                .maxSize(maxSize)
                .maxWaitMillis(maxWaitMillis)
                .connectionTimeoutMillis(connectionTimeoutMillis)
                .validationQuery(validationQuery)
                .build()
                .getConnection();
    }

    // 关闭连接工具方法
    private void closeConnection(Connection connection) {
        if (connection != null) {
            try {
                if (!connection.isClosed()) {
                    connection.close();
                }
            } catch (SQLException e) {
                Logger.getLogger(SimpleDataSource.class.getName()).warning("Error closing connection: " + e.getMessage());
            }
        }
    }

    // 关闭所有连接
    public void close() {
        // 中断清理线程
        Thread.currentThread().interrupt();

        // 关闭所有连接
        PooledConnection conn;
        while ((conn = connectionPool.poll()) != null) {
            closeConnection(conn.connection);
        }

        // 清理资源
        connectionPool.clear();
    }

    // 以下方法为DataSource接口的默认实现
    @Override
    public PrintWriter getLogWriter() throws SQLException {
        return DriverManager.getLogWriter();
    }

    @Override
    public void setLogWriter(PrintWriter out) throws SQLException {
        DriverManager.setLogWriter(out);
    }

    @Override
    public void setLoginTimeout(int seconds) throws SQLException {
        DriverManager.setLoginTimeout(seconds);
    }

    @Override
    public int getLoginTimeout() throws SQLException {
        return DriverManager.getLoginTimeout();
    }

    @Override
    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
        throw new SQLFeatureNotSupportedException("getParentLogger not supported");
    }

    @Override
    public <T> T unwrap(Class<T> iface) throws SQLException {
        if (iface.isInstance(this)) {
            return (T) this;
        }
        throw new SQLException("DataSource does not support unwrapping to " + iface.getName());
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return iface.isInstance(this);
    }

    // 获取当前活跃连接数
    public synchronized int getActiveConnections() {
        return activeConnections;
    }

    // 获取当前空闲连接数
    public int getIdleConnections() {
        return connectionPool.size();
    }

    // 建造者模式，方便配置数据源参数
    public static class Builder {
        private String driverClass;
        private String url;
        private String username;
        private String password;
        private int initialSize = 5;
        private int maxSize = 10;
        private long maxWaitMillis = 30000; // 默认等待30秒
        private long connectionTimeoutMillis = 3600000; // 默认连接超时1小时
        private String validationQuery = "SELECT 1"; // 简单的验证查询

        public Builder driverClass(String driverClass) {
            this.driverClass = driverClass;
            return this;
        }

        public Builder url(String url) {
            this.url = url;
            return this;
        }

        public Builder username(String username) {
            this.username = username;
            return this;
        }

        public Builder password(String password) {
            this.password = password;
            return this;
        }

        public Builder initialSize(int initialSize) {
            this.initialSize = initialSize;
            return this;
        }

        public Builder maxSize(int maxSize) {
            this.maxSize = maxSize;
            return this;
        }

        public Builder maxWaitMillis(long maxWaitMillis) {
            this.maxWaitMillis = maxWaitMillis;
            return this;
        }

        public Builder connectionTimeoutMillis(long connectionTimeoutMillis) {
            this.connectionTimeoutMillis = connectionTimeoutMillis;
            return this;
        }

        public Builder validationQuery(String validationQuery) {
            this.validationQuery = validationQuery;
            return this;
        }

        public SimpleDataSource build() {
            if (driverClass == null || url == null) {
                throw new IllegalArgumentException("Driver class and URL must be set");
            }
            return new SimpleDataSource(this);
        }
    }
    public static Builder builder() {
        return new Builder();
    }
}
